from factor.BaseFactor import ArrayManagerFactorBase
import talib
import numpy as np
from typing import Union


class ArrayManagerFactor(ArrayManagerFactorBase):
    def __init__(self, size, symbol, interval):
        super().__init__(size=size, symbol=symbol, interval=interval)

    def sma(self, n: int = 5, array: bool = False) -> Union[float, np.ndarray]:
        """
        Simple moving average.
        """
        result: np.ndarray = talib.SMA(self.close, n)
        if array:
            return result
        return result[-1]

    def sma7(self, n: int = 7, array: bool = False) -> Union[float, np.ndarray]:
        """
        Simple moving average.
        """
        result: np.ndarray = talib.SMA(self.close, n)
        if array:
            return result
        return result[-1]

    def sma21(self, n: int = 21, array: bool = False) -> Union[float, np.ndarray]:
        """
        Simple moving average.
        """
        result: np.ndarray = talib.SMA(self.close, n)
        if array:
            return result
        return result[-1]

    def sma60(self, n: int = 60, array: bool = False) -> Union[float, np.ndarray]:
        """
        Simple moving average.
        """
        result: np.ndarray = talib.SMA(self.close, n)
        if array:
            return result
        return result[-1]

    def max21(self, n: int = 21, array: bool = False) -> Union[float, np.ndarray]:
        """
        Max value of period.
        """
        result: np.ndarray = np.zeros(self.size)
        for i in range(self.size):
            if i >= n-1:
                result[i] = np.max(self.close[i-n+1:i+1])
        if array:
            return result
        return result[-1]

    def min21(self, n: int = 21, array: bool = False) -> Union[float, np.ndarray]:
        """
        Min value of period.
        """
        result: np.ndarray = np.zeros(self.size)
        for i in range(self.size):
            if i >= n-1:
                result[i] = np.min(self.close[i-n+1:i+1])
        if array:
            return result
        return result[-1]

    def dd(self, n: int = 0, array: bool = False) -> Union[float, np.ndarray]:
        """
        自定义因子dd，因子计算方法为后一个K线的close减去前一个K线的open，差为因子值
        """
        # 初始化数据维度与self.close一致的result数据
        result: np.ndarray = np.zeros(len(self.close))
        # 覆盖result中对应位置的数据为因子值
        result[1:] = self.close[1:] - self.open[:1]
        if array:
            return result
        return result[-1]





